package io.xccit.deepseek.utils;

import io.xccit.deepseek.enums.MessageType;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;

/**
 * @author CH_ywx
 * @version 1.0
 * @create 2025/4/7
 * <p>SseServer</p>
 */
@Slf4j
public class SseServer {

    public static final Map<String, SseEmitter> sseClients = new ConcurrentHashMap<>();

    /**
     * 在线人数
     */
    private static AtomicInteger count = new AtomicInteger(0);

    /**
     * 建立SSE连接
     * @return
     */
    public static SseEmitter sseConnection(String userId){
        SseEmitter sseEmitter = new SseEmitter(0L);
        sseEmitter.onCompletion(completionCallback(userId));
        sseEmitter.onError(errorCallback(userId));
        sseEmitter.onTimeout(timeoutCallback(userId));
        try {
            sseEmitter.send(SseEmitter.event().name("connect").data("SSE连接成功"));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        sseClients.put(userId,sseEmitter);
        count.getAndIncrement();
        log.info("sse连接建立成功,用户ID为{}",userId);
        return sseEmitter;
    }

    /**
     * 发送sse消息
     * @param userId 用户ID
     * @param message 消息
     * @param messageType 消息类型
     */
    public static void sendMessage(String userId, String message, MessageType messageType) {
        SseEmitter sseEmitter = sseClients.get(userId);
        if (sseEmitter != null) {
            try {
                sendMessage(sseEmitter,userId,message,messageType);
            } catch (Exception e) {
                sseClients.remove(userId);
            }
        }
    }

    /**
     * 发送sse消息
     * @param sseEmitter sseEmitter
     * @param userId 用户ID
     * @param message 消息
     * @param messageType 消息类型
     */
    private static void sendMessage(SseEmitter sseEmitter, String userId, String message, MessageType messageType) {
        if (sseEmitter != null){
            try {
                sseEmitter.send(SseEmitter.event()
                        .name(messageType.getType())
                        .data(message));
            } catch (Exception e) {
                sseClients.remove(userId);
                log.error("sse消息发送失败,用户ID为{}",userId);
            }
        }
    }

    /**
     * 广播sse消息
     * @param message
     */
    public static void broadCast(String message){
        sseClients.forEach((userId, sseEmitter) -> {
            sendMessage(sseEmitter,userId,message,MessageType.MESSAGE);
        });
    }

    /**
     * 获取sse连接数
     * @return
     */
    public static int getCount() {
        return count.get();
    }


    /**
     * 系统主动断开sse连接
     * @param userId
     */
    public static void stopConnected(String userId){
        SseEmitter sseEmitter = sseClients.get(userId);
        if (sseEmitter != null){
            sseEmitter.complete();
            sseClients.remove(userId);
            count.getAndDecrement();
            log.info("sse连接关闭,用户ID为{}",userId);
        }
    }

    /**
     * 超时回调
     * @param userId
     * @return
     */
    private static Runnable timeoutCallback(String userId) {
        return () -> {
            sseClients.remove(userId);
            log.info("sse连接超时,用户ID为{}",userId);
        };
    }

    /**
     * 错误回调
     * @param userId
     * @return
     */
    private static Consumer<Throwable> errorCallback(String userId) {
        return throwable -> {
            sseClients.remove(userId);
            count.getAndDecrement();
            log.error("sse连接发生异常,用户ID为{}",userId);
        };
    }

    /**
     * 事件完成回调
     * @param userId
     * @return
     */
    private static Runnable completionCallback(String userId) {
        return () -> {
            sseClients.remove(userId);
            count.getAndDecrement();
            log.info("sse连接完成,用户ID为{}",userId);
        };
    }

    /**
     * 关闭sse连接
     * @param userId
     */
    public static void close(String userId){
        SseEmitter sseEmitter = sseClients.get(userId);
        if (sseEmitter != null){
            sseEmitter.complete();
            sseClients.remove(userId);
            count.getAndDecrement();
            log.info("sse连接关闭,用户ID为{}",userId);
        }
    }
}
